import os

import cv2
import matplotlib
import pandas as pd
import matplotlib.pyplot as plt
from moviepy.audio.io.AudioFileClip import AudioFileClip
from moviepy.video.compositing.CompositeVideoClip import CompositeVideoClip
from moviepy.video.io.VideoFileClip import VideoFileClip

# 设置Matplotlib以支持中文显示和正常显示负号
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为黑体
matplotlib.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# (1) 处理数据, (文件路径)
def preprocess_data(file_path):
    # 读取Excel文件
    data = pd.read_excel(file_path, index_col=0)
    # 将年份索引转换为整数
    data.index = data.index.astype(int)
    # 生成1900年到2024年的日期范围，以月为单位（只包含年月）
    monthly_index = pd.date_range(start='1900-01', end='2024-12', freq='MS').strftime('%Y-%m')
    # 创建一个空的DataFrame，用于存储扩展后的数据
    expanded_data = pd.DataFrame(index=monthly_index, columns=data.columns)

    # 将原始数据按年份对应到 1 月份
    for year in data.index:
        expanded_data.loc[str(year) + '-01'] = data.loc[year].values

    # 确保所有列为数值类型
    expanded_data = expanded_data.apply(pd.to_numeric, errors='coerce')

    # 设定初始值为0
    for column in expanded_data.columns:
        # 找到每列最前面的正整数索引
        first_nonzero_idx = expanded_data[column][expanded_data[column] > 0].index.min()
        if pd.notna(first_nonzero_idx):  # 确保first_nonzero_idx有效
            # 将最前面的正整数之前的所有值设为0（不包括该索引）
            expanded_data.loc[expanded_data.index < first_nonzero_idx, column] = 0

    # 对现有数据进行线性插值，axis=0表示按列插值
    interpolated_data = expanded_data.interpolate(method='linear', axis=0)

    return interpolated_data

# (2) 画图, (处理后的数据，保存路径)
def draw(data, path, watermark):
    if not os.path.exists(path):
        os.makedirs(path)
    # 自定义颜色
    colors = {
        '中国': 'red',
        '英国': 'blue',
        '巴西': 'green',
        '法国': 'purple',
        '德国': 'brown',
        '西班牙': 'orange',
        '俄罗斯': 'pink',
        '日本': 'gray',
        '印度': 'yellow',
        '韩国': 'cyan',
        '美国': 'magenta'
    }
    tag = "1900"
    cnt = 1
    for i in data.index:
        print("正在画图 {}".format(cnt))
        cnt += 1
        data_selected_date = data.loc[i]
        # 按值排序，并取前10个国家
        top_countries = data_selected_date.sort_values(ascending=True).tail(10)

        # 确保所有国家都有颜色定义
        all_countries = data.columns.tolist()
        for country in all_countries:
            if country not in colors:
                colors[country] = 'black'

        # 绘制横向的柱状图
        fig, ax = plt.subplots(figsize=(10, 6))
        bars = ax.barh(top_countries.index, top_countries.values,
                       color=[colors[country] for country in top_countries.index])
        # 设置国家名称在条形内部的右侧或外部的左侧
        # 添加柱状图标签
        for bar, country in zip(bars, top_countries.index):
            ax.text(bar.get_width(), bar.get_y() + bar.get_height() / 2,
                    f'{bar.get_width():.2f}km', va='center', ha='left', color='black')
            if bar.get_width() != 0:
                ax.text(bar.get_width() / 2, bar.get_y() + bar.get_height() / 2,
                               country, va='center', ha='center', color='black')

        ax.set_xlim(0, max(top_countries.values) * 1.1)  # 增加一点空间
        ax.set_title('全球各国地铁里程')
        if int(i.split('-')[0]) % 5 == 0:
            tag = i.split('-')[0]
        ax.text(0.9, -0.06, tag, transform=ax.transAxes, va='top', ha='right', color='black',
                fontsize=16)

        alpha_value = 0.5  # 透明度值（0.0为完全透明，1.0为完全不透明）
        x_position, y_position = 0.95, 0.05
        ax.text(x_position, y_position, watermark, fontsize=14, alpha=alpha_value, color='gray', ha='right',
                va='bottom',
                transform=ax.transAxes)
        # 保存图片
        plt.savefig(f'{path}/{i}.png')
        plt.close()

# (3) 生成视频, (图片路径)
def creat_video(path):
    # 创建存储图片的目录
    os.makedirs(path, exist_ok=True)

    # 设置视频参数
    frame_size = (1600, 1200)  # 调整为匹配保存的图像大小
    video_filename = 'vout.mp4'
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 编解码器
    fps = 12

    # 初始化视频写入对象
    video_writer = cv2.VideoWriter(video_filename, fourcc, fps, frame_size)

    # 读取图片并写入视频
    img_array = []
    ls = sorted(os.listdir('pic'))  # 确保按顺序读取文件

    for filename in ls:
        img_path = os.path.join('pic', filename)
        img = cv2.imread(img_path)

        # 确保图片尺寸与视频帧尺寸一致
        if img.shape[1] != frame_size[0] or img.shape[0] != frame_size[1]:
            img = cv2.resize(img, frame_size)

        img_array.append(img)

    for img in img_array:
        video_writer.write(img)

    video_writer.release()
    print('视频生成完成，请查看文件:', video_filename)

# (4) 添加音频 (视频路径, 音频路径), 找个时长大于视频的音频
def add_bgm(video_filename, audio_filename):
    # 读取视频和音频文件
    video_clip = VideoFileClip(video_filename)
    audio_clip = AudioFileClip(audio_filename)

    # 确保音频的长度和视频的长度相同，或者根据你的需求裁剪音频
    # 如果音频比视频长，你可以使用set_duration(video_clip.duration)来裁剪音频
    audio_clip = audio_clip.set_duration(video_clip.duration)

    # 创建一个包含视频和音频的CompositeVideoClip对象
    final_clip = CompositeVideoClip([video_clip.set_audio(audio_clip)])

    # 写入新的视频文件，包含原始视频和新的音频
    final_clip.write_videofile("vout_with_bgm.mp4", codec='libx264', audio_codec='aac')

    # 关闭所有剪辑
    video_clip.close()
    audio_clip.close()
    final_clip.close()

    print('带有背景音乐的视频生成完成，请查看文件:', "vout_with_bgm.mp4")

# 处理的xlsx
file_path = '国家地铁里程历程.xlsx'
data = preprocess_data(file_path)

# 保存过程图片目录
output_path = "pic"
watermark = ''  # 水印文本
draw(data, output_path, watermark)
creat_video(output_path)

# 视频文件名和音频文件名
# video_filename = 'vout.mp4'  # 带有水印的视频文件名
# audio_filename = 'bgm.mp3'  # 背景音乐文件名
# add_bgm(video_filename, audio_filename)

